from sklearn.datasets import load_digits  # 导入sklearn的数据集，用于加载手写数字数据集
from sklearn.model_selection import train_test_split  # 导入数据集划分工具，用于将数据集划分为训练集和测试集
from sklearn.neighbors import KNeighborsClassifier  # 导入K近邻分类器，用于进行K近邻分类
import matplotlib.pyplot as plt  # 导入matplotlib的pyplot模块，用于绘制图像
import pickle  # 导入pickle模块，用于保存和加载模型
from tqdm import tqdm  # 导入tqdm模块，用于显示进度条


# 加载手写数字数据集
digits = load_digits()

# 将数据集划分为训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(
    digits.data, digits.target, test_size=0.2, random_state=42)

# 初始化变量，用于存储最佳准确率和对应的k值
best_accuracy = 0
best_k = 0
best_knn = None

# 初始化列表，用于存储每个k值的准确率
accuracy_list = []

# 尝试k值从1到40
for k in tqdm(range(1, 41)):
    # 使用当前k值初始化K近邻分类器
    knn = KNeighborsClassifier(n_neighbors=k)
    
    # 在训练数据上拟合模型
    knn.fit(X_train, y_train)
    
    # 计算模型的准确率
    accuracy = knn.score(X_test, y_test)
    
    # 将准确率添加到列表中
    accuracy_list.append(accuracy)
    
    # 更新最佳准确率和对应的k值
    if accuracy > best_accuracy:
        best_accuracy = accuracy
        best_k = k
        best_knn = knn

# 输出最佳准确率和对应的k值
print(f"最佳准确率为 {best_accuracy:.2f}，对应的k值为 {best_k}")

# 将最佳KNN模型保存到二进制文件
with open('best_knn_model.pkl', 'wb') as file:
    pickle.dump(best_knn, file)


# 绘制每个k值的准确率
plt.plot(range(1, 41), accuracy_list)
plt.xlabel('k value')
plt.ylabel('Accuracy')
plt.title('Accuracy of different k values', fontsize=15)

# Draw a red vertical line at the k value with the highest accuracy
plt.axvline(x=best_k, color='red')

# Mark the point with the highest accuracy and annotate the k value and accuracy
plt.scatter(best_k, best_accuracy, color='red')  # Mark the highest accuracy
plt.text(best_k, best_accuracy, f'k={best_k}, Accuracy={best_accuracy:.2f}', color='red')

# 将图像保存为PDF文件
plt.savefig('accuracy_plot.pdf', format='pdf')